#include "crope.h"

#include <buola/graph/cgraph.h>
#include <buola/graph/components.h>
#include <buola/graph/dijkstra.h>
#include <buola/graph/csubgraph.h>
#include <buola/graph/floyd.h>
#include <buola/cv/ccloud.h>
#include <buola/scene/geodes/cplane.h>
#include <buola/scene/cscene.h>
#include <buola/iterator/transform.h>

#include <BulletDynamics/ConstraintSolver/btPoint2PointConstraint.h>
#include <BulletDynamics/ConstraintSolver/btGeneric6DofSpringConstraint.h>

namespace buola { namespace tracker {

static const float sKDRope=0;
static const float sKPRope=0.2;

static const float sScale=100;
static const float sGraphProximity=0.04f*sScale;
static const float sNodeLength=0.02f*sScale;
static const float sBinSize=0.06f*sScale;
static const float sRopeRadius=0.006f*sScale;
static const float sNodeMass=0.2;
//for rope and rope2
static const float sPlanePos=-1.08f*sScale;
//for sequence2
//static const float sPlanePos=-1.58f*sScale;
static const int sDownSample=1;

struct ARopeFilter
{
    bool operator()(const cvi::CCloud_rgb8::TPoint &pP) const
    {
        //for rope!
        if(pP.mPixel[0]>50||pP.mPixel[1]>50||pP.mPixel[2]<50) return false;
        if(pP.mPixel[2]<=pP.mPixel[1]+10||pP.mPixel[2]<=pP.mPixel[0]+10) return false;

        //for rope2!
//        if(pP.mPixel[2]<=pP.mPixel[1]||pP.mPixel[2]<=pP.mPixel[0]) return false;


        if(pP.mPoint.x()<-0.45f*sScale) return false;
        if(pP.mPoint.x()>0.3f*sScale) return false;
        if(pP.mPoint.y()<0.6f*sScale) return false;
        if(pP.mPoint.y()>2.0f*sScale) return false;
//         if(pP.mPoint.z()>-0.7) return false;
        
        
        
        return true;
    }
};
    
typedef typename cvi::CCloudPoint_rgb8 TRopeNode;
    
template<typename tGraph>
graph::CGraph<TRopeNode,float> calc_reeb_graph(const tGraph &pGraph)
{
    auto lComp=graph::connected_components(pGraph);
    auto lDistance=graph::new_vertex_property<float>(pGraph);
    auto lBinIndex=graph::new_vertex_property<int>(pGraph);

    std::vector<std::vector<typename tGraph::vertex_descriptor>> lPartitions;
    
    for(auto &c : lComp)
    {
        sort(c);
        graph::dijkstra_shortest_distance(pGraph,c[0],pGraph.EdgeMap(),lDistance.Map());
        
        transform(lDistance[c],lBinIndex[c],[](float d){return (int)(d/sBinSize);});

        int lMaxBinIndex=max_value(lBinIndex[c]);

        for(int i=0;i<=lMaxBinIndex;i++)
        {
            std::vector<typename tGraph::vertex_descriptor> lVertices;
            for(auto v : c)
            {
                if(lBinIndex[v]==i)
                    lVertices.push_back(v);
            }
            
            auto lSubGraph=graph::make_subgraph(pGraph,lVertices);
            auto lSubComp=graph::connected_components(lSubGraph);

            lPartitions.insert(lPartitions.end(),move_begin(lSubComp),move_end(lSubComp));
        }
    }
    
    graph::CGraph<TRopeNode,float> lSkeleton;
    auto lPartIndex=graph::new_vertex_property<graph::CGraph<mat::CVec3f,float>::vertex_descriptor>(pGraph);
    for(const auto &p : lPartitions)
    {
        mat::CVec3f lCenter=accumulate(pGraph[p]|transform([](const TRopeNode &n){return n.mPoint;}),mat::CVec3f(0,0,0))/p.size(); //find center of partition
        img::CColor lColor=accumulate(pGraph[p]|transform([](const TRopeNode &n){return img::CColor(0,0,0);}),img::CColor(0,0,0,0))/p.size(); //find average color TODO
        TRopeNode lNode;
        lNode.mPoint=lCenter;
        lNode.mPixel=lColor;
        auto lNewVertex=boost::add_vertex(lNode,lSkeleton); //add it to skeleton
        for(auto v : p) //mark every vertex in partition as belonging to that point in skeleton
            lPartIndex[v]=lNewVertex;
    }

    for(int i=0;i<lPartitions.size();i++)
    {
        for(int j=i+1;j<lPartitions.size();j++)
        {
            bool lDone=false;
            for(auto v : lPartitions[i])
            {
                for(auto e : range(out_edges(v,pGraph)))
                {
                    if(lPartIndex[target(e,pGraph)]==j)
                    {
                        add_edge(i,j,mat::distance2(lSkeleton[i].mPoint,lSkeleton[j].mPoint),lSkeleton);
                        lDone=true;
                        break;
                    }
                }
                if(lDone) break;
            }
        }
    }
    
    return lSkeleton;
}

template<typename tGraph>
std::vector<std::vector<typename tGraph::vertex_descriptor>> get_segments(const tGraph &pGraph)
{
    std::vector<std::vector<typename tGraph::vertex_descriptor>> lSegments;
    
    std::vector<typename tGraph::vertex_descriptor> lJunctions;

    for(auto v : vertices(pGraph))
    {
        if(out_degree(v,pGraph)!=2)
            lJunctions.push_back(v);
    }
    
    for(auto j : lJunctions)
    {
        for(auto e : range(out_edges(j,pGraph)))
        {
            std::vector<typename tGraph::vertex_descriptor> lSegment;
            
            auto l1=j;
            auto l2=target(e,pGraph);

            lSegment.push_back(l1);
            lSegment.push_back(l2);
            
            while(out_degree(l2,pGraph)==2)
            {
                for(auto e : range(out_edges(l2,pGraph)))
                {
                    if(target(e,pGraph)==l1) continue;
                    l1=l2;
                    l2=target(e,pGraph);
                    lSegment.push_back(l2);
                    break;
                }
            }
            
            if(lSegment.front()>lSegment.back()) continue;
            
            if(lSegment.front()==lSegment.back())
            {
                bool lFound=false;
                for(auto s : lSegments)
                {
                    if(s.size()!=lSegment.size()) continue;
                    if(equals(s,range(lSegment.rbegin(),lSegment.rend())))
                    {
                        lFound=true;
                        break;
                    }
                }
                
                if(lFound) continue;
            }

            lSegments.push_back(std::move(lSegment));
        }
    }
    
    return lSegments;
}

graph::CGraph<TRopeNode> skeletonize_point_cloud(const cvi::CCloud_rgb8 &pCloud)
{
    std::vector<TRopeNode> lPoints;
    for(const auto &p : pCloud)
    {
        TRopeNode lNode;
        lPoints.emplace_back(p);
    }
    
    auto lGraph=graph::build(lPoints,[](const TRopeNode &p1,const TRopeNode &p2){return distance(p1.mPoint,p2.mPoint)<sGraphProximity;},
                                     [](const TRopeNode &p1,const TRopeNode &p2){return distance(p1.mPoint,p2.mPoint);});
    auto lSkeleton=calc_reeb_graph(lGraph);
     
    auto lPair=graph::longest_shortest_path(lSkeleton,lSkeleton.EdgeMap());
    auto lPath=graph::dijkstra_shortest_path(lSkeleton,lPair.first,lPair.second,lSkeleton.EdgeMap());
    graph::CGraph<TRopeNode> lPruned;
    for(int i=0;i<lPath.size();i++)
    {
        add_vertex(lSkeleton[lPath[i]],lPruned);
        if(i>0)
            add_edge(num_vertices(lPruned)-2,num_vertices(lPruned)-1,lPruned);
    }
    
    return lPruned;
}

std::vector<TRopeNode> smooth(const std::vector<TRopeNode> &pSrc,float pLength)
{
    int N=pSrc.size();
    int M=3; //number of components
    float lTotalDiff=0;
    mat::CVec_f lX(N);
    mat::CVec_f lDiffX(N-1);
    std::vector<mat::CRow_f> lDiffY;
    
    for(int i=0;i<N-1;i++)
    {
        lX[i]=lTotalDiff;
        lDiffX[i]=distance(pSrc[i].mPoint,pSrc[i+1].mPoint);
        lTotalDiff+=lDiffX[i];
        lDiffY.push_back((pSrc[i+1].mPoint-pSrc[i].mPoint).T());
    }
    lX(N-1)=INFINITY;
    
    mat::CMat_f lL=mat::zeros(N,3);
    mat::CMat_f lR=mat::zeros(N,M);
    
    //build tridiagonal matrix
    lL(0,1)=2.0f/(lDiffX[0]);
    lL(0,2)=1.0f/(lDiffX[0]);
    lR(0,nAll)=3.0f*(lDiffY[0])/sq(lDiffX[0]);
    for(int i=1;i<N-1;i++)
    {
        lL(i,0)=1.0f/lDiffX[i-1];
        lL(i,1)=2.0f/lDiffX[i-1]+2.0f/lDiffX[i];
        lL(i,2)=1.0f/lDiffX[i];
        lR(i,nAll)=3.0f*(lDiffY[i-1]/sq(lDiffX[i-1])+lDiffY[i]/sq(lDiffX[i]));
    }
    lL(N-1,0)=1.0f/lDiffX[N-2];
    lL(N-1,1)=2.0f/lDiffX[N-2];
    lR(N-1,nAll)=3.0f*lDiffY[N-2]/sq(lDiffX[N-2]);;
    
    //forward pass
    lL(0,2)=lL(0,2)/lL(0,1);
    lR(0,nAll)=lR(0,nAll)/lL(0,1);
    for(int i=1;i<N-1;i++)
    {
        lL(i,2)=lL(i,2)/(lL(i,1)-lL(i-1,2)*lL(i,0));
        lR(i,nAll)=(lR(i,nAll)-lR(i-1,nAll)*lL(i,0))/(lL(i,1)-lL(i-1,2)*lL(i,0));
    }
    
    for(int i=N-2;i>=0;i--)
    {
        lR(i,nAll)=lR(i,nAll)-lL(i,2)*lR(i+1,nAll);
    }
    
    std::vector<TRopeNode> lResult;
    int lNumSegs=(int)round(lTotalDiff/pLength);
    pLength=lTotalDiff/lNumSegs;
    float lCumLength=0;
    int lPos;
    for(int i=0;i<=lNumSegs;i++)
    {
        while(lCumLength>lX[lPos+1]) lPos++;
        
        mat::CVec3f lY2=pSrc[lPos+1].mPoint;
        mat::CVec3f lY1=pSrc[lPos].mPoint;
        mat::CRow_f a=lR(lPos,nAll)*lDiffX[lPos]-lDiffY[lPos];
        mat::CRow_f b=-lR(lPos+1,nAll)*lDiffX[lPos]+lDiffY[lPos];
        float t=(lCumLength-lX[lPos])/lDiffX[lPos];
        TRopeNode lY;
        lY.mPixel=pSrc[lPos].mPixel;
        lY.mPoint=((1-t)*lY1+t*lY2+t*(1-t)*(a*(1-t)+b*t).T());
        msg_info() << lY.mPoint.T() << "\n";
        lResult.push_back(lY);
        lCumLength+=pLength;
    }
    return lResult;
}

CRope::CRope(const cvi::CCloud_rgb8 &pCloud)
{
    auto lFiltered=cvi::filter(pCloud,ARopeFilter());
//    auto lFiltered=clone(pCloud);

    auto lSkeleton=skeletonize_point_cloud(lFiltered);
    auto lNodes=make_vector(lSkeleton[vertices(lSkeleton)]);
    lNodes=smooth(lNodes,sNodeLength);

    float lLastLength;
    for(int i=1;i<lNodes.size();i++)
    {
        mat::CVec3f lCenter=(lNodes[i].mPoint+lNodes[i-1].mPoint)/2;
        mat::CVec3f lDiff=lNodes[i].mPoint-lNodes[i-1].mPoint;
        float lLength=norm(lDiff);
        mat::CQuat_f lRotation(mat::CVec3f(0,1,0),lDiff);
        mat::CVec3f lPerp=mat::cross(lDiff,mat::CVec3f(0,0,1));
        lPerp=mat::cross(lPerp,lDiff);
        mat::CQuat_f lRotation2(lRotation*mat::CVec3f(0,0,1),lPerp);
        scene::PGeode lGeode=new scene::CCapsule(lLength,sRopeRadius);
//            lGeode->SetColor(lNodes[i].mColor);
        lGeode->SetColor(img::nice_color((i-1)/5));
        scene::PRigidBody lBody=new scene::CRigidBody(lGeode,sNodeMass,
                                                      mat::CMat34d(mat::rotation_translation_3f(lRotation2*lRotation,lCenter)));
        
        msg_info() << lLength << "\n";
        
        if(i>=2)
        {
            auto *lP2P=new btPoint2PointConstraint(*mLinks.back()->GetBulletBody(),*lBody->GetBulletBody(),
                                                            btVector3(0,lLastLength/2,0),btVector3(0,-lLength/2,0));
            lP2P->setParam(BT_CONSTRAINT_STOP_ERP,0.2);
            mConstraints.push_back(lP2P);
            
            btTransform lA;
            lA.setIdentity();
            lA.setOrigin(btVector3(0,lLastLength/2,0));
            btTransform lB;
            lB.setIdentity();
            lB.setOrigin(btVector3(0,-lLength/2,0));
            auto *lSpring=new btGeneric6DofSpringConstraint(*mLinks.back()->GetBulletBody(),*lBody->GetBulletBody(),
                                                            lA,lB,false);
            for(int i=3;i<=5;i++)
            {
                lSpring->enableSpring(i,true);
                lSpring->setStiffness(i,0.1);
                lSpring->setDamping(i,1);
            }
            
            float lLimit=0.1;
            lSpring->setAngularLowerLimit(btVector3(-lLimit,-lLimit,-lLimit));
            lSpring->setAngularUpperLimit(btVector3(lLimit,lLimit,lLimit));
            
            mConstraints.push_back(lSpring);
        }

        mLinks.push_back(lBody);
        lLastLength=lLength;
    }
    
    mPlane=new scene::CRigidBody(new scene::CPlane,0,mat::translation_3d({0,0,sPlanePos}));
}

int CRope::GetK()
{
    return mLinks.size();
}

cvi::CCloud_rgb8 CRope::FilterCloud(const cvi::CCloud_rgb8& pCloud)
{
    return cvi::filter(pCloud,ARopeFilter());
}

mat::CMat_f CRope::GetCloudFeatures(const cvi::CCloud_rgb8& pCloud)
{
    mat::CMat_f lResult(pCloud.size()/sDownSample,3);

    msg_info() << "N:" << pCloud.size() << "\n";
    
    for(int i=0;i<pCloud.size()/sDownSample;i++)
    {
        lResult(i,nAll)=pCloud[i*sDownSample].mPoint.T();
    }
    
    return lResult;
}

mat::CMat_f CRope::GetObjectFeatures()
{
    mat::CMat_f lResult(mLinks.size(),3);
    
    msg_info() << "K:" << mLinks.size() << "\n";
    
    for(int i=0;i<mLinks.size();i++)
    {
        lResult(i,nAll)=mat::CVec3f(mLinks[i]->CenterOfMass()).T();
    }
    
    return lResult;
}

void CRope::AddToScene(const scene::PScene &pScene)
{
    for(const auto &b : mLinks)
    {
        pScene->World()->AddChild(b);
        pScene->AddBody(b);
    }
    
    pScene->World()->AddChild(mPlane);
    pScene->AddBody(mPlane);
    
    for(const auto &c : mConstraints)
    {
        pScene->AddConstraint(c);
    }
}

void CRope::ApplyEvidence(const mat::CMat_f &pPZGivenC,const mat::CMat_f &pObserved,const mat::CMat_f &pObject)
{
    for(int k=0;k<mLinks.size();k++) 
    { 
        scene::PRigidBody lBody=mLinks[k];
        btVector3 lImpulse=-sKDRope*lBody->GetBulletBody()->getLinearVelocity();
        for(int n=0;n<pPZGivenC.Cols();n++)
        {
            btVector3 lObsPt(pObserved(n,0),pObserved(n,1),pObserved(n,2));
            btVector3 lEstPt(pObject(k,0),pObject(k,1),pObject(k,2));
            lImpulse+=(sKPRope*pPZGivenC(k,n))*(lObsPt-lEstPt);
        }
        
        lImpulse=lImpulse/lBody->GetBulletBody()->getInvMass();
        if(k%10==0)
            msg_info() << "impulse " << k << " " <<lImpulse[0] << " " << lImpulse[1] << " " << lImpulse[2] << "\n";
        lBody->GetBulletBody()->applyCentralImpulse(lImpulse);
    }
    end_timer("M step");
}

/*namespace tracker*/ } /*namespace buola*/ }
